from Runtime.utils import *
from Runtime.elf_config import e_machine_dit


class ELF_Templates(object):
    address = []
    pass


class ELF_Error(Exception):
    def __init__(self, error_info):
        super().__init__(self)
        self.error_info = error_info

    def __str__(self):
        return self.error_info


class ELF_Header(ELF_Templates):
    elf_header, elf_header_e_ident = dict(), dict()
    elf_bit, elf_endian, elf_version, elf_type, elf_machine, elf_os_abi = r'', r'', 255, r'', r'', 0
    elf_bytes, elf_io = None, None
    elf_entry_addr, elf_program_header_start_offset, elf_section_header_start_offset = 0, 0, 0
    elf_flags, elf_header_size, elf_program_tab_size, elf_program_tab_number = 0, 0, 0, 0
    elf_section_tab_size, elf_section_tab_number, elf_section_name_string_table = 0, 0, 0
    elf_string_table_index = 0
    elf_header_item_size = [16, 2, 2, 4, 8, 8, 8, 4, 2, 2, 2, 2, 2, 2]

    def __init__(self, elf_bytes):
        self.elf_bytes = elf_bytes
        self.elf_header[r'e_ident'] = self.elf_bytes[0:15 + 1]
        self.elf_header[r'e_type'] = bytes_to_int(self.elf_bytes[16:17 + 1])
        self.elf_header[r'e_machine'] = e_machine_dit[bytes_to_int(self.elf_bytes[18:19 + 1])]
        self.elf_header[r'e_version'] = bytes_to_int(self.elf_bytes[20:23 + 1])
        self.elf_header[r'e_entry'] = bytes_to_int(self.elf_bytes[24:31 + 1])
        self.elf_header[r'e_phoff'] = bytes_to_int(self.elf_bytes[32:39 + 1], )
        self.elf_header[r'e_shoff'] = bytes_to_int(self.elf_bytes[40:47 + 1])
        self.elf_header[r'e_flags'] = self.elf_bytes[48:51 + 1]
        self.elf_header[r'e_ehsize'] = bytes_to_int(self.elf_bytes[52:53 + 1])
        self.elf_header[r'e_phentsize'] = bytes_to_int(self.elf_bytes[54:55 + 1])
        self.elf_header[r'e_phnum'] = bytes_to_int(self.elf_bytes[56:57 + 1])
        self.elf_header[r'e_shentsize'] = bytes_to_int(self.elf_bytes[58:59 + 1])
        self.elf_header[r'e_shnum'] = bytes_to_int(self.elf_bytes[60:61 + 1])
        self.elf_header[r'e_shstrndx'] = bytes_to_int(self.elf_bytes[62:63 + 1])

        self.elf_header_e_ident['EI_MAG0'] = chr(self.elf_bytes[0])
        self.elf_header_e_ident['EI_MAG1'] = chr(self.elf_bytes[1])
        self.elf_header_e_ident['EI_MAG2'] = chr(self.elf_bytes[2])
        self.elf_header_e_ident['EI_MAG3'] = chr(self.elf_bytes[3])
        self.elf_header_e_ident['EI_CLASS'] = int(self.elf_bytes[4])
        self.elf_header_e_ident['EI_DATA'] = int(self.elf_bytes[5])
        self.elf_header_e_ident['EI_VERSION'] = self.elf_bytes[6]
        self.elf_header_e_ident['EI_OSABI'] = int(self.elf_bytes[7])
        self.elf_header_e_ident['EI_ABIVERSION'] = int(self.elf_bytes[8])
        self.elf_header_e_ident['EI_NIDENT'] = self.elf_bytes[8:15 + 1]

        if self.elf_header[r'e_ident'][:4] != bytes([0x7f, 0x45, 0x4c, 0x46]):
            raise ELF_Error(r'read header_e_ident error')
        if self.elf_header_e_ident['EI_CLASS'] == 0:
            raise ELF_Error(r'invalid EI_CLASS')
        elif self.elf_header_e_ident['EI_CLASS'] == 1:
            raise ELF_Error(r'ELFCLASS32 not support!')
        elif self.elf_header_e_ident['EI_CLASS'] == 2:
            self.elf_bit = r'64'

        if self.elf_header_e_ident['EI_DATA'] != 1:
            raise ELF_Error(r'endian EI_DATA not support!')
        else:
            self.elf_endian = r'ELFDATA2LSB'

        self.elf_os_abi = self.elf_header_e_ident['EI_OSABI']
        self.elf_version = int(self.elf_header_e_ident['EI_VERSION'])

        if self.elf_header[r'e_type'] == 0 or self.elf_header[r'e_type'] == 0xff00 or self.elf_header[
            r'e_type'] == 0xffff:
            raise ELF_Error(r' e_type not support!')
        elif self.elf_header[r'e_type'] == 1:
            self.elf_type = r'Relocatable file'
        elif self.elf_header[r'e_type'] == 2:
            self.elf_type = r'Executable file'
        elif self.elf_header[r'e_type'] == 3:
            self.elf_type = r'Shared object file'
        elif self.elf_header[r'e_type'] == 4:
            self.elf_type = r'Core file'

        self.elf_machine = self.elf_header[r'e_machine']
        if self.elf_header[r'e_version'] != 1:
            raise ELF_Error(r'e_version not support!')

        self.elf_entry_addr = self.elf_header[r'e_entry']
        self.elf_program_header_start_offset = self.elf_header[r'e_phoff']
        self.elf_section_header_start_offset = self.elf_header[r'e_shoff']
        self.elf_flags = self.elf_header[r'e_flags']
        self.elf_header_size = self.elf_header[r'e_ehsize']
        self.elf_program_tab_size = self.elf_header[r'e_phentsize']  # 单个程序表的大小
        self.elf_program_tab_number = self.elf_header[r'e_phnum']  # 程序表的个数
        self.elf_section_tab_size = self.elf_header[r'e_shentsize']
        self.elf_section_tab_number = self.elf_header[r'e_shnum']
        self.elf_string_table_index = self.elf_header[r'e_shstrndx']

    def __len__(self):
        ret = 0
        for i in self.elf_header_item_size:
            ret += i
        return ret


class ELF_Program_Header(ELF_Templates):
    look_up_list = [
        ("p_type", 4, "segment_type"),
        ("p_flags", 4, "segment_flags"),
        ("p_offset", 8, "segment_file_offset"),
        ("p_vaddr", 8, "segment_virtual_address"),
        ("p_paddr", 8, "segment_physical_address"),
        ("p_filesz", 8, "segment_size_in_file"),
        ("p_memsz", 8, "segment_size_in_memory"),
        ("p_align", 8, "segment_alignment"),
    ]
    _bytes = None
    segment_type = None
    segment_flags = None
    segment_file_offset = None
    segment_virtual_address = None
    segment_physical_address = None
    segment_size_in_file = None
    segment_size_in_memory = None
    segment_alignment = None

    def __init__(self, elf_bytes, start):
        self._bytes = elf_bytes
        offset = start
        for p_name, step, attr in self.look_up_list:
            setattr(self, attr, bytes_to_int(self._bytes[offset:offset + step]))
            self.address.append((offset, offset + step - 1))
            offset += step

    def __len__(self):
        ret = 0
        for _, i, _ in self.look_up_list:
            ret += i
        return ret


class EFL_Section_Header(ELF_Templates):
    look_up_list = [
        ("sh_name", 4, "Section_name"),
        ("sh_type", 4, "Section_type"),
        ("sh_flags", 8, "Section_flags"),
        ("sh_addr", 8, "Section_virtual_addr"),
        ("sh_offset", 8, "Section_file_offset"),
        ("sh_size", 8, "Section_size_in_bytes"),
        ("sh_link", 4, "Section_Link_to_another_section"),
        ("sh_info", 4, "Section_Additional_section_information"),
        ("sh_addralign", 8, "Section_alignment"),
        ("sh_entsize", 8, "Section_Entry_size_if_section_holds_table"),
    ]
    def __init__(self, elf_bytes, start):
        self.Section_name = None
        self.Section_type = None
        self.Section_flags = None
        self.Section_virtual_addr = None
        self.Section_file_offset = None
        self.Section_size_in_bytes = None
        self.Link_to_another_section = None
        self.Additional_section_information = None
        self.Section_alignment = None
        self.Entry_size_if_section_holds_table = None
        self._bytes = elf_bytes
        offset = start
        for p_name, step, attr in self.look_up_list:
            # print(offset, offset + step)
            setattr(self, attr, bytes_to_int(self._bytes[offset:offset + step]))
            self.address.append((offset, offset + step - 1))
            offset += step

    def __len__(self):
        ret = 0
        for _, i, _ in self.look_up_list:
            ret += i
        return ret

class ELF(ELF_Templates):
    elf_header = None
    elf_program_headers = []  # 这里装的都是segment的头 信息表
    elf_section_headers = []

    def __init__(self, file_name):
        with open(file_name, "rb") as elf_file_io:
            self.elf_bytes = elf_file_io.read()
            self.elf_io = elf_file_io
        self.elf_header = ELF_Header(self.elf_bytes)
        program_start = self.elf_header.elf_program_header_start_offset
        for _ in range(self.elf_header.elf_program_tab_number):
            self.elf_program_headers.append(ELF_Program_Header(self.elf_bytes, program_start))
            program_start += self.elf_header.elf_program_tab_size

        if self.elf_header.elf_section_tab_number == 0:
            raise ELF_Error(r'this elf file section table = 0 !')
        section_start = self.elf_header.elf_section_header_start_offset
        for _ in range(self.elf_header.elf_section_tab_number):
            self.elf_section_headers.append(EFL_Section_Header(self.elf_bytes, section_start))
            section_start += self.elf_header.elf_section_tab_size
